import tensorflow as tf
import numpy as np

class AdaptiveEarlyStopping(tf.keras.callbacks.Callback):
    """自适应早停回调"""
    
    def __init__(self, patience=10, min_delta=0.001, restore_best_weights=True):
        super(AdaptiveEarlyStopping, self).__init__()
        self.patience = patience
        self.min_delta = min_delta
        self.restore_best_weights = restore_best_weights
        self.best_weights = None
        self.best_loss = np.Inf
        self.wait = 0
    
    def on_train_begin(self, logs=None):
        self.wait = 0
        self.best_loss = np.Inf
    
    def on_epoch_end(self, epoch, logs=None):
        current_loss = logs.get('val_loss')
        if current_loss is None:
            return
        
        if current_loss < self.best_loss - self.min_delta:
            self.best_loss = current_loss
            self.wait = 0
            if self.restore_best_weights:
                self.best_weights = self.model.get_weights()
        else:
            self.wait += 1
            if self.wait >= self.patience:
                self.model.stop_training = True
                if self.restore_best_weights and self.best_weights is not None:
                    self.model.set_weights(self.best_weights)